use std::slice;
use std::env;
use std::sync::Mutex;
use ::async_nats;
use ::tokio;
use tokio::time::Duration;
use tokio::runtime::Runtime;
use crate::event::ExternalSequenceBlockHash;
mod event;


lazy_static::lazy_static! {
    // 全局Tokio运行时
    static ref g_rt: Runtime = tokio::runtime::Runtime::new().unwrap();

    // 全局配置字符串
    static ref KV_EVENT_SUBJECT: String = {
        // 从环境变量读取配置
        let namespace = env::var("VLLM_KV_NAMESPACE")
            .unwrap_or_else(|_| "dynamo".to_string());

        let component = env::var("VLLM_KV_COMPONENT")
            .unwrap_or_else(|_| "VllmWorker".to_string());

        // 格式化字符串
        format!("namespace.{}.component.{}.kv_events", namespace, component)
    };

    // Worker id
    static ref g_worker_id: i64 = {
        env::var("VLLM_WORKER_ID")
            .and_then(|s| s.parse::<i64>().map_err(|_| env::VarError::NotPresent))
            .unwrap_or(0)
    };

    // 全局NATS客户端
    static ref NATS_CLIENT: Mutex<Option<async_nats::Client>> = Mutex::new(None);
}

#[unsafe(no_mangle)]
pub extern "C" fn create_nats_connection() -> i32 {
    let nats_server = env::var("NATS_SERVER")
        .unwrap_or_else(|_| "nats://localhost:4222".to_string());

    // 在同步上下文中执行异步连接
    let client = g_rt.block_on(async {
        let options = async_nats::ConnectOptions::with_user_and_password("user".to_string(), "user".to_string())
            .ping_interval(Duration::from_secs(30))  // 每30秒发送心跳
            .max_reconnects(Some(30))               // 最多重试30次
            .reconnect_delay_callback(|attempts| {
                Duration::from_secs(std::cmp::min(attempts as u64 * 2, 30)) // 显式转换为u64
            })
            .event_callback(|event| async move {
                println!("NATS Connect Event: {:?}", event);
            });

        options.connect(nats_server).await
    });
    match client {
        Ok(client) => {
            println!("Connected to NATS!");
            *NATS_CLIENT.lock().unwrap() = Some(client);
            0
        },
        Err(e) => {
            eprintln!("Failed to connect to NATS: {}", e);
            1
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn free_nats_connection() {
    let mut client = NATS_CLIENT.lock().unwrap();
    if client.take().is_some() {
        println!("NATS connection released successfully");
    } else {
        eprintln!("No NATS connection to release");
    }
}

#[derive(Clone)]
#[repr(C)]
pub struct StoredBlockData {
    external_hash: u64,
    local_hash: u64,
}

#[unsafe(no_mangle)]
pub extern "C" fn publishStoredEvent(event_id: u64, parent_hash: u64, blk_info: *const StoredBlockData, len: usize, is_first: bool) -> i32 {
    if *g_worker_id == 0 {
        eprintln!("invalid work id: 0, please check env variables: VLLM_WORKER_ID");
        return 0;
    }

    assert!(!blk_info.is_null(), "Null pointer received");
    let blk_vec = unsafe { slice::from_raw_parts(blk_info, len) }.to_vec();
    let store_data = event::KvCacheStoreData {
        blocks: blk_vec
            .iter()
            .map(|blk_data| event::KvCacheStoredBlockData {
                block_hash: event::ExternalSequenceBlockHash(blk_data.external_hash),
                tokens_hash: event::LocalBlockHash(blk_data.local_hash),
            })
            .collect(),
        parent_hash: if is_first {
            None
        } else {
            Some(event::ExternalSequenceBlockHash(parent_hash))
        },
    };
    let data = event::KvCacheEventData::Stored(store_data);
    let event = event::KvCacheEvent { event_id, data };
    let router_event = event::RouterEvent { worker_id: *g_worker_id, event };

    let ret = {
        if let Err(e) = event::publish_event(&router_event) {
            eprintln!("publish failed: {}", e);
            1
        } else {
            println!("publish stored event finish!");
            0
        }
    };

    ret
}

#[unsafe(no_mangle)]
pub extern "C" fn publishRemovedEvent(event_id: u64, hash_ids: *const u64, len: usize) -> i32 {
    assert!(!hash_ids.is_null(), "Null pointer received");
    let blk_vec: Vec<ExternalSequenceBlockHash> = unsafe { slice::from_raw_parts(hash_ids, len) }
        .to_vec()
        .iter()
        .map(|&v| ExternalSequenceBlockHash(v))
        .collect();
    let remove_data = event::KvCacheRemoveData {
        block_hashes: blk_vec,
    };
    let data         = event::KvCacheEventData::Removed(remove_data);
    let event        = event::KvCacheEvent { event_id, data };
    let router_event = event::RouterEvent { worker_id: *g_worker_id, event };

    let ret = {
        if let Err(e) = event::publish_event(&router_event) {
            eprintln!("publish failed: {}", e);
            1
        } else {
            println!("publish removed event finish!");
            0
        }
    };

    ret
}
